/*
 * Copyright (C) 2022 Open Source Robotics Foundation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
*/

#include <gz/msgs/param.pb.h>
#include <gz/msgs/vector3d.pb.h>
#include <memory>
#include <string>
#include <gz/sensors/Noise.hh>
#include <gz/sim/components/Pose.hh>
#include <gz/sim/Conversions.hh>
#include <gz/sim/Entity.hh>
#include <gz/sim/Util.hh>
#include <gz/plugin/Register.hh>
#include <gz/transport/Node.hh>
#include <sdf/sdf.hh>

#include "AcousticPingerPlugin.hh"

using namespace gz;
using namespace vrx;

/// \brief Private AcousticPingerPlugin data class.
class AcousticPingerPlugin::Implementation
{
  /// \brief Parse all SDF parameters.
  /// \param[in] _ecm The Entity Component Manager.
  /// \return True when all params were successfully parsed or false otherwise.
  public: bool ParseSDFParameters(sim::EntityComponentManager &_ecm);

  /// \brief Callback function called when receiving a new pinger position
  /// via the pinger subscription callback.
  /// \param[in] _pos New pinger position.
  public: void PingerPositionCallback(const msgs::Vector3d &_pos);

  /// \brief Transport node.
  public: transport::Node node;

  /// \brief Publisher used to send sensor messages generated by the plugin.
  public: transport::Node::Publisher rangeBearingPub;

  /// \brief Mutex to protect the position vector.
  public: std::mutex mutex;

  /// \brief Entity of the model attached to the plugin.
  public: sim::Entity entity = sim::kNullEntity;

  /// \brief Vector storing the position of the pinger.
  public: math::Vector3d position = math::Vector3d::Zero;

  /// \brief String holding the frame id of the sensor.
  public: std::string frameId = "pinger";

  /// \brief The topic used to read the range bearing measurements.
  public: std::string topic = "/pinger/range_bearing";

  /// \brief The topic to set the position of the pinger.
  public: std::string setPositionTopicName = "/pinger/set_pinger_position";

  /// \brief Sensor update rate.
  public: float updateRate = 1.0f;

  /// \brief The next message to be published.
  public: msgs::Param msg;

  /// \brief SDF pointer.
  public: sdf::ElementPtr sdf;

  /// \brief Variable used to track time of last update. This is used to
  /// produce data at the correct rate.
  public: std::chrono::duration<double> lastUpdateTime{0};

  // From Brian Bingham's rangebearing_gazebo_plugin.
  /// \brief rangeNoise - Gazebo noise object for range.
  public: sensors::NoisePtr rangeNoise = nullptr;

  /// \brief Gazebo noise object for bearing angle.
  public: sensors::NoisePtr bearingNoise = nullptr;

  /// \brief Gazebo noise object for elevation angle.
  public: sensors::NoisePtr elevationNoise = nullptr;
};

//////////////////////////////////////////////////
bool AcousticPingerPlugin::Implementation::ParseSDFParameters(
  sim::EntityComponentManager &_ecm)
{
  // Set the frame_id. Defaults to "pinger".
  if (this->sdf->HasElement("frame_id"))
    this->frameId = this->sdf->Get<std::string>("frame_id");

  // Load topic from SDF if available.
  if (this->sdf->HasElement("topic"))
    this->topic = this->sdf->Get<std::string>("topic");

  // Set the topic to be used to publish the sensor message.
  if (this->sdf->HasElement("set_position_topic"))
  {
    this->setPositionTopicName =
      this->sdf->Get<std::string>("set_position_topic");
  }

  // Initialize pinger position. Defaults to origin.
  if (this->sdf->HasElement("position"))
    this->position = this->sdf->Get<math::Vector3d>("position");

  // Initialise update rate. Default to 1 reading per second.
  if (this->sdf->HasElement("update_rate"))
    this->updateRate = this->sdf->Get<float>("update_rate");

  // From Brian Bingham's rangebearing_gazebo_plugin.
  // Noise setup and parse SDF.
  if (this->sdf->HasElement("range_noise"))
  {
    sdf::ElementPtr rangeNoiseElem = this->sdf->GetElement("range_noise");
    // Note - it is hardcoded into the NoiseFactory.cc that the SDF
    // element be "noise".
    if (rangeNoiseElem->HasElement("noise"))
    {
      this->rangeNoise = sensors::NoiseFactory::NewNoiseModel(
        rangeNoiseElem->GetElement("noise"));
    }
    else
    {
      gzerr << "<range_noise> must contain a noise tag" << std::endl;
    }
  }

  // Load the noise model from the SDF file.
  if (this->sdf->HasElement("bearing_noise"))
  {
    sdf::ElementPtr bearingNoiseElem = this->sdf->GetElement("bearing_noise");
    // Note - it is hardcoded into the NoiseFactory.cc that the SDF
    // element be "noise".
    if (bearingNoiseElem->HasElement("noise"))
    {
      this->bearingNoise = sensors::NoiseFactory::NewNoiseModel(
        bearingNoiseElem->GetElement("noise"));
    }
    else
    {
      gzerr << "<bearing_noise> must contain a noise tag" << std::endl;
    }
  }

  if (this->sdf->HasElement("elevation_noise"))
  {
    sdf::ElementPtr elevationNoiseElem =
      this->sdf->GetElement("elevation_noise");
    // Note - it is hardcoded into the NoiseFactory.cc that the SDF
    // element be "noise".
    if (elevationNoiseElem->HasElement("noise"))
    {
      this->elevationNoise = sensors::NoiseFactory::NewNoiseModel(
        elevationNoiseElem->GetElement("noise"));
    }
    else
    {
      gzerr << "<elevation_noise> must contain a noise tag" << std::endl;
    }
  }

  return true;
}

//////////////////////////////////////////////////
void AcousticPingerPlugin::Implementation::PingerPositionCallback(
  const msgs::Vector3d &_pos)
{
  std::lock_guard<std::mutex> lock(this->mutex);
  this->position = math::Vector3d(_pos.x(), _pos.y(), _pos.z());
}

//////////////////////////////////////////////////
AcousticPingerPlugin::AcousticPingerPlugin()
  : dataPtr(utils::MakeUniqueImpl<Implementation>())
{
}

//////////////////////////////////////////////////
void AcousticPingerPlugin::Configure(const sim::Entity &_entity,
  const std::shared_ptr<const sdf::Element> &_sdf,
  sim::EntityComponentManager &_ecm, sim::EventManager &_eventMgr)
{
  this->dataPtr->sdf = _sdf->Clone();
  if (!this->dataPtr->ParseSDFParameters(_ecm))
  {
    gzerr << "AcousticPingerPlugin disabled" << std::endl;
    return;
  }

  // Setup the publisher.
  this->dataPtr->rangeBearingPub =
    this->dataPtr->node.Advertise<msgs::Param>(this->dataPtr->topic);

  // Subscriber to update the pinger pose.
  this->dataPtr->node.Subscribe(this->dataPtr->setPositionTopicName,
    &AcousticPingerPlugin::Implementation::PingerPositionCallback,
    this->dataPtr.get());

  this->dataPtr->entity = _entity;

  // Prepopulate the msg.
  auto *param = this->dataPtr->msg.mutable_params();

  msgs::Any rangeValue;
  rangeValue.set_type(msgs::Any_ValueType::Any_ValueType_DOUBLE);
  rangeValue.set_double_value(0.0);
  (*param)["range"] = rangeValue;

  msgs::Any bearingValue;
  bearingValue.set_type(msgs::Any_ValueType::Any_ValueType_DOUBLE);
  bearingValue.set_double_value(0.0);
  (*param)["bearing"] = bearingValue;

  msgs::Any elevationValue;
  elevationValue.set_type(msgs::Any_ValueType::Any_ValueType_DOUBLE);
  elevationValue.set_double_value(0.0);
  (*param)["elevation"] = elevationValue;
}

//////////////////////////////////////////////////
void AcousticPingerPlugin::PreUpdate(const sim::UpdateInfo &_info,
  sim::EntityComponentManager &_ecm)
{
  if (_info.paused)
    return;

  if ((_info.simTime - this->dataPtr->lastUpdateTime) <
      (std::chrono::duration<double>(1.0) / this->dataPtr->updateRate))
  {
    return;
  }

  this->dataPtr->lastUpdateTime = _info.simTime;

  auto modelPose =
    _ecm.Component<sim::components::Pose>(this->dataPtr->entity)->Data();

  math::Vector3d direction;
  {
    std::lock_guard<std::mutex> lock(this->dataPtr->mutex);

    // Direction vector to the pinger from the USV.
    direction = this->dataPtr->position - modelPose.Pos();
  }

  // Sensor reading is in the sensor frame. Rotate the direction vector into
  // the frame of reference of the sensor.
  math::Vector3d directionSensorFrame =
    modelPose.Rot().RotateVectorReverse(direction);

  // Generate a 2D vector for elevation calculation.
  math::Vector3d directionSensorFrame2d =
    math::Vector3d(directionSensorFrame.X(), directionSensorFrame.Y(), 0);

  // Bearing is calculated by finding the world frame direction vector
  // and transforming into the pose of the sensor. Bearing is calculated
  // using the atan2 function of the x and y components of the transformed
  // vector. The elevation is calculated from the length of the 2D only
  // and the z component of the sensor frame vector.
  double bearing = atan2(directionSensorFrame.Y(), directionSensorFrame.X());
  double range = directionSensorFrame.Length();
  double elevation =
    atan2(directionSensorFrame.Z(), directionSensorFrame2d.Length());

  // Apply noise to each measurement.
  // From Brian Binghams rangebearing_gazebo_plugin.
  if (this->dataPtr->rangeNoise)
    range = this->dataPtr->rangeNoise->Apply(range);
  if (this->dataPtr->bearingNoise)
    bearing = this->dataPtr->bearingNoise->Apply(bearing);
  if (this->dataPtr->elevationNoise)
    elevation  = this->dataPtr->elevationNoise->Apply(elevation);

  // Update message
  //   - stamp
  auto stamp = sim::convert<msgs::Time>(_info.simTime);
  this->dataPtr->msg.mutable_header()->mutable_stamp()->CopyFrom(stamp);

  //   - frame_id
  auto frame = this->dataPtr->msg.mutable_header()->add_data();
  frame->set_key("frame_id");
  frame->add_value(this->dataPtr->frameId);

  //   - range, bearing and elevation
  auto *param = this->dataPtr->msg.mutable_params();
  (*param)["range"].set_double_value(range);
  (*param)["bearing"].set_double_value(bearing);
  (*param)["elevation"].set_double_value(elevation);

  // Publish a new measurement.
  this->dataPtr->rangeBearingPub.Publish(this->dataPtr->msg);
}

GZ_ADD_PLUGIN(AcousticPingerPlugin,
              sim::System,
              AcousticPingerPlugin::ISystemConfigure,
              AcousticPingerPlugin::ISystemPreUpdate)

GZ_ADD_PLUGIN_ALIAS(vrx::AcousticPingerPlugin,
                    "vrx::AcousticPingerPlugin")
